#这是用共轭梯度法求解大规模稀疏方程组的程序

import numpy as np
import pandas as pd
import warnings

warnings.filterwarnings("error")#警告转错误，避免0/0
#共轭梯度法
class CGM():
    def __init__(self, A: np.ndarray, b: np.ndarray, x=[], e=0.5*10**(-2), pot=16):
        """
        初始化共轭梯度法！
        :param A: 创建实例时需传入Ax=b的矩阵A与向量b,numpy类型；
        :param b: 创建实例时需传入Ax=b的矩阵A与向量b,numpy类型；
        :param x: 初始向量x可以被修改，默认x(0)=[0,0,0……],列表类型；
        :param e: 精度e默认为0.5*10^(-2),可以修改；
        :param pot:小数点位数，默认16位，最大16位。
        """
        self.pot =pot
        self.A = A
        self.b = b
        if x != []:
            try:
                self.x = np.array([x]).T
            except:
                self.x = np.zeros((A.shape[1], 1))
        else:
            self.x = np.zeros((A.shape[1], 1))
        self.e = e
        self.flag = True #求解完成的标志

    def forward(self):
        """
        这是CGM的计算过程。
        :return:x 或 error
        """
        print("你的输入如下：")
        print("x:", self.x)
        print("A:", self.A)
        print("b", self.b)
        print("e", self.e)
        try:
            d = self.b - self.A @ self.x
            r = self.b - self.A @ self.x
            if (np.abs(r) < self.e).all():
                print("求解结果：", self.x)
                print("-*-" * 20)
                return np.around(self.x, self.pot)
        except:
            print("数据输入有误！请检查！")
            print("-*-" * 20)
            return "数据输入有误！请检查！"
        n = 0
        while 1:
            try:
                a = r.T @ r / (d.T @ self.A @ d)
                self.x = self.x + a * d
                r0 = r
                r = self.b - self.A @ self.x
                if np.max(np.abs(r)) < self.e or n + 1 == self.A.shape[0]:
                    break
                # B = -r.T @ self.A @ d / (d.T @ self.A @ d)
                # B = r.T @ self.A @ r / (d.T @ self.A @ d)
                B = (((r**2).sum())**(1/2)) / (((r0**2).sum())**(1/2))
                d = r + B * d
            except:
                print("无解，或初始值不能让迭代收敛，请检查相关输入！")
                print("-*-" * 20)
                self.flag = False
                return "无解，或初始值不能让迭代收敛，请检查相关输入！"
            n += 1
        if self.flag == True:
            print("求解结果：", self.x)
            print("-*-" * 20)
            return np.around(self.x, self.pot)

    def params_change(self, **kwargs):
        """
        这是实例生成后还需要修改相关参数时调用的函数。
        :param kwargs: 对需要修改的参数，以 A= ,b= ,x= ,e= 的形式传入，其数据类型与建立实例时一样。
        :return: 无返回值
        """
        try:
            self.A = kwargs["A"]
        except:
            self.A = self.A
        try:
            self.b = kwargs["b"]
        except:
            self.b = self.b
        try:
            self.x = np.array(kwargs["x"]).T
        except:
            if self.x.shape[0] != self.A.shape[0]:
                self.x = np.zeros((A.shape[1], 1))
            else:
                self.x = self.x
        try:
            self.e = kwargs["e"]
        except:
            self.e = self.e

#3.2题的矩阵生成函数
def Ab(n):
    """
    生成系数矩阵A和值b的方法。
    :param n: 系数矩阵A的阶数。
    :return:A,b
    """
    b = np.zeros((n, 1))
    b[0][0] = -1
    b[n-1][0] = -1
    A = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            if i==j:
                A[i][j] = -2
            elif i == j+1 or j == i+1:
                A[i][j] = 1
    return A, b

def save_ans(path, ans):
    """
    保存结果的函数。
    :param path: 文件名及路径,str类型
    :param ans: 传入结果。
    :return: 无
    """
    ans = pd.DataFrame(ans, columns=["结果"])
    ans.to_csv(path + ".csv", index=False)

if __name__ == '__main__':
    #调试用范例
    # cgm1 = CGM(A=np.ones((5, 5)), b=np.ones((5, 1)), x=[1,1,1,1,1])
    # ans1 = cgm1.forward()

    #
    # A1 = np.array([[2,0,1],[0,1,0],[1,0,2]])
    # b1 = np.array([[3],[1],[3]])
    # cgm2 = CGM(A=A1, b=b1, x=[1, 1, 1])
    # ans2 = cgm2.forward()

    # A2 = np.array([[4,1],[1,3]])
    # b2 = np.array([[1],[2]])
    # cgm3 = CGM(A=A2, b=b2, x=[0, 0])
    # ans3 = cgm3.forward()

    #计算实习3.2
    # print("A的阶数为100时的结果:")
    # A, b = Ab(100)
    # cgm3_2 = CGM(A=A, b=b)
    # ans3_2_100 = cgm3_2.forward()
    # save_ans("./data/共轭梯度法3_2_100", ans3_2_100)
    #
    # print("A的阶数为200时的结果:")
    # A, b = Ab(200)
    # cgm3_2.params_change(A=A, b=b)
    # ans3_2_200 = cgm3_2.forward()
    # save_ans("./data/共轭梯度法3_2_200", ans3_2_200)
    #
    # print("A的阶数为400时的结果:")
    # A, b = Ab(500)
    # # x = [[1 for i in range(2000)]]
    # cgm3_2.params_change(A=A, b=b)
    # ans3_2_400 = cgm3_2.forward()
    # save_ans("./data/共轭梯度法3_2_400", ans3_2_400)
    #
    # print("A的阶数为400时的结果（加大精度后）:")
    # A, b = Ab(400)
    # cgm3_2.params_change(A=A, b=b, e=0.5*10**(-3))
    # ans3_2_400 = cgm3_2.forward()
    # save_ans("./data/共轭梯度法3_2_400(高精度)", ans3_2_400)

    #磁悬浮
    A = pd.read_csv("./data/K.csv", index_col=0, header=0)
    b = pd.read_csv("./data/F.csv", index_col=0, header=0)
    A = np.array(A)
    b = np.array(b)
    cxf = CGM(A=A, b=b)
    cxf.forward()
    save_ans("./data/磁悬浮求解结果", cxf.forward())
    print("直接矩阵计算结果：\n", np.linalg.inv(A) @ b)